#include "yolov5_ascend.h"

static const cv::Scalar CV_COLOR_BLUE(255, 0, 0);
static const cv::Scalar CV_COLOR_GREEN(0, 255, 0);
static const cv::Scalar CV_COLOR_RED(0, 0, 255);
static const cv::Scalar CV_COLOR_BLACK(0, 0 ,0);
static const cv::Scalar CV_COLOR_WHITE(255, 255, 255);

static bool objectCompare(const DetectionTarget& obj1, const DetectionTarget& obj2)
{
    return obj1.conf > obj2.conf;
}

static void nmSortedboxes(const std::vector<DetectionTarget>& objects,
    std::vector<int>& picked, float nmsThr)
{
    picked.clear();
    const int n = objects.size();
    std::vector<float> areas(n);
    for (int i=0; i<n; ++i)
    {
        const float& x1 = objects[i].x1;
        const float& y1 = objects[i].y1;
        const float& x2 = objects[i].x2;
        const float& y2 = objects[i].y2;
        areas[i] = fabs(x2-x1) * fabs(y2-y1);
    }

    // NMS with all bboxes (across classes)
    cv::Rect2f rect1, rect2;
    for (int i=0; i < n; ++i)
    {
        const DetectionTarget& det1 = objects[i];
        rect1 = cv::Rect2f(det1.x1, det1.y1, det1.x2-det1.x1, det1.y2-det1.y1);
        
        int keep = 1;
        for (int j=0; j<(int)picked.size(); ++j)
        {
            const DetectionTarget& det2 = objects[picked[j]];
            rect2 = cv::Rect2f(det2.x1, det2.y1, det2.x2-det2.x1, det2.y2-det2.y1);
            float iArea = (rect1 & rect2).area();
            float uArea = areas[i] + areas[picked[j]] - iArea;
            if (iArea / uArea > nmsThr)
            {
                keep = 0;
                break;
            }
        }
        if (keep)
        {
            picked.push_back(i);
        }
    }
}

Yolov5Ascend::Yolov5Ascend(const Yolov5AscendDetConfigs& config):modelId(0), ascendCfg(config)
{
    uint32_t deviceNum;
    auto ret =aclrtGetDeviceCount(&deviceNum);
    if (ret != 0)
    {
        std::cout << "acl get device count failed" << std::endl;
        std::exit(EXIT_FAILURE);
    }
    std::cout << "available device num: " << deviceNum << std::endl;
    if (config.deviceId >= (int)deviceNum)
    {
        std::cout << "invalid device id" << std::endl;
        std::exit(EXIT_FAILURE);
    }

    this->pictureDataSize = ascendCfg.inputSize.area() * 3;
    init();
}

Yolov5Ascend::~Yolov5Ascend()
{
    this->unloadModel();
    this->unloadPicture();
    this->destroyResource();
}

void Yolov5Ascend::init()
{
    if(!this->initResource())
    {
        return;
    }
    if(!this->loadModelFromFile(ascendCfg.omPath.c_str()))
    {
        return;
    }
    if(!this->createDesc())
    {
        return;
    }
    aclError temp = aclrtMalloc(&this->pictureDeviceData, this->pictureDataSize, ACL_MEM_MALLOC_HUGE_FIRST);
    if (temp!=ACL_SUCCESS)
    {
        std::cout << "apply device data failed"<< std::endl;
        return;
    }
    this->createModelInput();
    this->createModelOutput();
}

bool Yolov5Ascend::initResource()
{
    aclError ret = aclInit(nullptr);
    if (ret!=ACL_SUCCESS)
    {
        std::cout << "aclInit error" << std::endl;
        return false;
    } 

    ret = aclrtSetDevice(this->ascendCfg.deviceId);
    if (ret != ACL_SUCCESS)
    {
        std::cout<<"acl open device" << this->ascendCfg.deviceId <<" failed"<< std::endl;
        return false;
    }

    return true;
}

bool Yolov5Ascend::loadModelFromFile(const char *modelPath)
{
    // load model and get modelID.
    aclError ret = aclmdlLoadFromFile(modelPath, &this->modelId);
    if (ret != ACL_SUCCESS)
    {
        std::cout << "load model from file failed, model file is "<< modelPath<< std::endl;
        return false;
    }

    return true;
}

bool Yolov5Ascend::createDesc()
{
    this->modelDesc = aclmdlCreateDesc();
    if (this->modelDesc == nullptr)
    {
        std::cout << "create model description failed" << std::endl;
        return false;
    }
    // get modelDesc(model description) by modelID
    aclError ret = aclmdlGetDesc(this->modelDesc, this->modelId);
    if (ret != ACL_SUCCESS)
    {
        std::cout << "get model description failed" << std::endl;
        return false;
    }

    return true;
}

void Yolov5Ascend::createModelInput()
{
    // 创建aclmdlDataset类型的数据，描述模型推理的输入
    this->inputDataSet = aclmdlCreateDataset();
    this->inputDataBuffer = aclCreateDataBuffer(this->pictureDeviceData, this->pictureDataSize);
    aclError ret = aclmdlAddDatasetBuffer(this->inputDataSet, this->inputDataBuffer);
    if (ret != 0)
    {
        std::cout << "add dataset buffer failed" << std::endl;
        std::exit(EXIT_FAILURE);
    }
}

void Yolov5Ascend::createModelOutput()
{
    // 创建aclmdlDataset类型的数据，描述模型推理的输出
    this->outputDataSet = aclmdlCreateDataset();
    // 获取模型输出数据需占用的内存大小，单位为Byte
    this->outputDataSize = aclmdlGetOutputSizeByIndex(this->modelDesc, 0);
    // 申请输出内存
    aclError ret = aclrtMalloc(&this->outputDeviceData, this->outputDataSize, ACL_MEM_MALLOC_HUGE_FIRST);
    if (ret != 0)
    {
        std::cout << "apply device data failed" << std::endl;
        std::exit(EXIT_FAILURE);
    }
    this->outputDataBuffer = aclCreateDataBuffer(this->outputDeviceData, this->outputDataSize);
    ret = aclmdlAddDatasetBuffer(this->outputDataSet, this->outputDataBuffer);
}

void Yolov5Ascend::unloadModel()
{
    aclError ret = aclmdlUnload(this->modelId);
    if (ret != ACL_SUCCESS)
    {
        std::cout << "unload model failed, modelId is " <<this->modelId<< std::endl;
    }

    if (this->modelDesc!= nullptr)
    {
        (void)aclmdlDestroyDesc(this->modelDesc);
        this->modelDesc= nullptr;
    }
}

void Yolov5Ascend::unloadPicture()
{
    aclError ret = aclrtFree(this->pictureDeviceData);
    if (ret != 0)
    {
        std::cout << "free device data failed" << std::endl;
        std::exit(EXIT_FAILURE);
    }
	this->pictureDeviceData = nullptr;
	aclDestroyDataBuffer(this->inputDataBuffer);
	this->inputDataBuffer = nullptr;
	aclmdlDestroyDataset(this->inputDataSet);
	this->inputDataSet = nullptr;

	ret = aclrtFree(this->outputDeviceData);
	this->outputDeviceData = nullptr;
	aclDestroyDataBuffer(this->outputDataBuffer);
	this->outputDataBuffer = nullptr;
	aclmdlDestroyDataset(this->outputDataSet);
	this->outputDataSet = nullptr;
}

void Yolov5Ascend::destroyResource()
{
	aclError ret = aclrtResetDevice(this->ascendCfg.deviceId);
    if (ret != 0)
    {
        std::cout << "reset device failed" << std::endl;
        std::exit(EXIT_FAILURE);
    }
	aclFinalize();
}

cv::Mat Yolov5Ascend::preProcess(cv::Mat& src, int interpolation)
{
    const int& inputW = ascendCfg.inputSize.width;
    const int& inputH = ascendCfg.inputSize.height;
    float r = std::min(inputW/(src.cols*1.0), inputH/(src.rows*1.0));
    int unpadW = r * src.cols;
    int unpadH = r * src.rows;
    cv::Mat re(unpadH, unpadW, CV_8UC3);
    cv::resize(src, re, re.size(), 0, 0, interpolation);
    cv::Mat padImg(inputH, inputW, CV_8UC3, cv::Scalar(114, 114, 114));
    re.copyTo(padImg(cv::Rect(0, 0, re.cols, re.rows)));
    // auto t0 = std::chrono::steady_clock::now();
    return padImg;
}

void Yolov5Ascend::postProcess(std::vector<DetectionTarget>& objects, const cv::Size& imgSize)
{
    void *outputHostData;
    aclrtMallocHost(&outputHostData, this->outputDataSize);
    aclrtMemcpy(outputHostData, this->outputDataSize, this->outputDeviceData, this->outputDataSize,
        ACL_MEMCPY_DEVICE_TO_HOST);
    float16_t *outFloatData = reinterpret_cast<float16_t *>(outputHostData);

    const int& classNum = ascendCfg.classNum;
    const int& inputW = ascendCfg.inputSize.width;
    const int& inputH = ascendCfg.inputSize.height;
    int proposalNum = (inputW/8) * (inputH/8) + (inputW/16) * (inputH/16) + (inputW/32) * (inputH/32);
    proposalNum *= 3; 
    std::vector<DetectionTarget> proposals;
    // auto t0 = std::chrono::steady_clock::now();
    for(int i = 0; i < proposalNum; i++)
    {
        float16_t *lineBegin = outFloatData + i*(classNum+5);
        float conf = float(lineBegin[4]);
        if (conf < ascendCfg.scoreThr || conf > 1.0)
        {
            continue;
        }
        int centerX = (int) (*(lineBegin));
        int centerY = (int) (*(lineBegin + 1));
        int width = (int) (*(lineBegin + 2));
        int height = (int) (*(lineBegin + 3));
        int maxClassId = 0;
        for (int k = 1; k < classNum; ++k)
        {
            if (lineBegin[5 + k] > lineBegin[5 + maxClassId])
            {
                maxClassId = k;
            }
        }
        float prob = conf * lineBegin[5+maxClassId];
        if (prob > ascendCfg.scoreThr)
        {
            DetectionTarget obj;
            obj.cls = maxClassId;
            obj.conf = prob;
            obj.typeName = "cls_" + std::to_string(obj.cls);
            obj.x1 = centerX - width / 2;
            obj.y1 = centerY - height / 2;
            obj.x2 = centerX + width / 2;
            obj.y2 = centerY + height / 2;
            proposals.push_back(obj);
        }
    }
    std::sort(proposals.begin(), proposals.end(), &objectCompare);
    std::vector<int> picked;
    nmSortedboxes(proposals, picked, ascendCfg.nmsThr);
    int count = picked.size();
    objects.resize(count);
    const int& imgW = imgSize.width;
    const int& imgH = imgSize.height;
    float scale = std::min(inputW / (imgW*1.0), inputH / (imgH*1.0));
    for (int i=0; i<count; ++i)
    {
        // scale
        objects[i] = proposals[picked[i]];
        objects[i].x1 /= scale;
        objects[i].y1 /= scale;
        objects[i].x2 /= scale;
        objects[i].y2 /= scale;
        // clip
        objects[i].x1 = std::max(std::min(objects[i].x1, (float)imgW-1), 0.0f);
        objects[i].y1 = std::max(std::min(objects[i].y1, (float)imgH-1), 0.0f);
        objects[i].x2 = std::max(std::min(objects[i].x2, (float)imgW-1), 0.0f);
        objects[i].y2 = std::max(std::min(objects[i].y2, (float)imgH-1), 0.0f);
    }
    aclrtFreeHost(outputHostData);
}

void Yolov5Ascend::detect(cv::Mat& img, std::vector<DetectionTarget>& objects)
{
    auto t0 = std::chrono::steady_clock::now();
    cv::Mat prImg = preProcess(img, cv::INTER_NEAREST);
    auto t1 = std::chrono::steady_clock::now();
    std::chrono::duration<double> elapsed = t1 - t0;
    float timeMs = elapsed.count() * 1000;
    // std::cout << "preprocess time " <<  timeMs << " ms\n";
    t0 = std::chrono::steady_clock::now();


    aclError ret = aclrtMemcpy(this->pictureDeviceData, this->pictureDataSize, prImg.ptr<uint8_t>(),
        this->pictureDataSize, ACL_MEMCPY_DEVICE_TO_DEVICE);
    if (ret!=ACL_SUCCESS)
    {
        std::cout << "copy data into the cache of the input dataset failed" << ret << std::endl;
        return;
    }
    ret = aclmdlExecute(this->modelId, this->inputDataSet, this->outputDataSet);
    if (ret!=ACL_SUCCESS)
    {
        std::cout << "model inference failed" << std::endl;
        return;
    }
    t1 = std::chrono::steady_clock::now();
    elapsed = t1 - t0;
    timeMs = elapsed.count() * 1000;
    // std::cout << "inference time " <<  timeMs << " ms\n";
    t0 = std::chrono::steady_clock::now();
    postProcess(objects, img.size());
    t1 = std::chrono::steady_clock::now();
    elapsed = t1 - t0;
    timeMs = elapsed.count() * 1000;
    // std::cout << "postprocess time " <<  timeMs << " ms\n";
}

void Yolov5Ascend::drawObject(cv::Mat& img, const std::vector<DetectionTarget>& objects,
    const std::vector<std::string>& labels)
{
    for (const auto& obj: objects)
    {
        cv::Rect rect = cv::Rect(obj.x1, obj.y1, obj.x2-obj.x1, obj.y2-obj.y1);
        cv::rectangle(img, rect, CV_COLOR_GREEN, 2);
        int baseLine = 0;
        int score = round(obj.conf * 100);
        std::string label = labels.size() == 0? obj.typeName: labels[obj.cls];
        label += ": " + std::to_string(score) + "%";
        cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX,
            1.0, 1, &baseLine);
        int x = obj.x1;
        int y = obj.y1 + 1;
        if (y > img.rows) 
        {
            y = img.rows;
        }
        cv::Rect backRect(cv::Point(x, y),
            cv::Size(labelSize.width, labelSize.height + baseLine));
        cv::rectangle(img, backRect, CV_COLOR_GREEN, -1);
        cv::putText(img, label.c_str(), cv::Point(x, y + labelSize.height),
            cv::FONT_HERSHEY_SIMPLEX, 1.0, CV_COLOR_RED, 2);
    }
}